clear; clc; close all
temp = dlmread('C:\Users\Mittal\Documents\Work\Code\survival-analysis\perl\normalizedEDPULSE.txt');
format short g

age = temp(:,2);
bp_rr = temp(:,3);
output = temp(:,4);
los = temp(:,5);

% Test inputs
% output = []; npts = 300;
% bp_rr = rand(npts,1)*3;
% knots = [1,2];
% output = double(bp_rr < knots(1) | bp_rr > knots(2));
% inds = unique(ceil(rand(25,1)*npts));
% output(inds) = abs(1-output(inds));

for nKnots = 1:25,
    dead = find(output == 1);
    alive = setdiff([1:length(output)],dead);
    
    figure(1); plot(bp_rr(alive),output(alive),'b.'); grid on; hold on;
    figure(1); plot(bp_rr(dead),output(dead),'r.'); grid on; hold on;
    xlabel('SYSTOLIC BP'); ylabel('PATIENT OUTCOME'); ylim([-1,3])
    
    % Count classwise occurances of each value of BP and put in Y.
    % Read description in mnrfit help
    % % x = unique(bp_rr);
    % % Y = nan(length(x),2);
    % % for i = 1:length(x),
    % %     temp = find(bp_rr == x(i));
    % %     Y(i,1) = length(find(output(temp) == 0));
    % %     Y(i,2) = length(find(output(temp) == 1));
    % % end
    % Alternatively
    x = bp_rr;
    x_sorted = sort(x);
    knots = nan(1,nKnots);
    for i = 1:nKnots,
        knots(i) = x_sorted(round(length(bp_rr)*i/(nKnots+1)));
        Ind = (bp_rr > knots(i));
        x = [x (bp_rr-knots(i)).*Ind];
    end
    Y = output; Y(find(Y == 0)) = 2;
    
    %Fit Logistic Regression
    [beta, ~, stat] = mnrfit(x,Y);
    % betaX = beta'*[ones(size(x,1),1) x]';
    
    npts = 1000;
    xx = linspace(min(bp_rr)-5,max(bp_rr)+5,npts)';
    X = [ones(npts,1) xx];
    
    for i = 1:length(knots),
        X = [X max(xx - knots(i),0)];
    end
    
    pHatNom = exp(beta'*X')./(1 + exp(beta'*X')); pHatNom = pHatNom';
    % pHatNom = mnrval(beta,xx,'model','nominal','interactions','on');
    plot(xx,pHatNom,'g'); if(nKnots > 0) vline(knots,'r'); end
    
    logLikelihood = (beta'*[ones(size(x,1),1) x]')*output - sum(log(1 + exp(beta'*[ones(size(x,1),1) x]')));
    AIC = length(beta)*2 - 2*logLikelihood;
    BIC = length(beta)*log(size(x,1)) - 2*logLikelihood;
    fprintf('knots = %d, AIC = %f, BIC = %f\n',nKnots,AIC,BIC);
    figure(2); plot(nKnots,AIC,'b*','markersize',6,'linewidth',4); hold on; plot(nKnots,BIC,'r+','markersize',6,'linewidth',4); grid on
    xlabel('Number of knots'); legend('AIC','BIC')
end

% [beta stat.p]